/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

// This demonstrates the SSL client-side support
//

package org.mozilla.jss.ssl;

import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.io.PrintStream;
import java.net.InetAddress;
import java.net.Socket;
import java.nio.charset.StandardCharsets;
import java.util.HashMap;

import org.mozilla.jss.CertDatabaseException;
import org.mozilla.jss.CryptoManager;
import org.mozilla.jss.InitializationValues;
import org.mozilla.jss.KeyDatabaseException;
import org.mozilla.jss.crypto.AlreadyInitializedException;

/**
 * Parameters supported by this socket test:
 *
 * filename file to be read from https server (default: /index.html)
 * port port to connect to (default: 443)
 * ipaddr address to connect to (overrides hostname, no default)
 * hostname host to connect to (no default)
 * clientauth do client-auth or not (default: no client-auth)
 *
 * The following parameters are used for regression testing, so
 * we can print success or failure of the test.
 *
 * filesize size of file to be read
 * status security status of connection - this has to be an integer
 * cipher
 * sessionKeySize
 * sessionSecretSize
 * issuer
 * subject
 * certSerialNum
 *
 */

public class SSLClient {
    boolean handshakeEventHappened = false;
    boolean doClientAuth = false;
    HashMap<String, String> args;
    PrintStream results;
    String versionStr;

    String[] argNames = {
            "filename",
            "port",
            "ipaddr",
            "hostname",
            "filesize",
            "status",
            "sessionKeySize",
            "sessionSecretSize",
            "cipher",
            "issuer",
            "subject",
            "certSerialNum",
    };

    String[] values = {
            "/index", // filename
            "443", // port to connect to
            "", // ipaddr (use hostname instead)
            "trading.etrade.com", // hostname to connect to
            "1024", // filesize
            "2", // status, 2 means ???
            "128", // expected session key size
            "128", // expected session key secret bits
            "RC48", // expected cipher
            "CN=Hardcore Certificate Server II, OU=Hardcore, O=Netscape Communications Corporation, C=US", // expected issuer DN of server
            "CN=hbombsgi.mcom.com, OU=Hardcore, C=US", // expected subject DN of server
            "00C3", // serial number
    };

    String okay = "okay";
    String failed = "FAILED";

    private static String htmlHeader = "SSL Client Tester";

    /* simple helper functions */

    private boolean isInvalid(String s) {
        return (s == null) || s.equals("");
    }

    private String getArgument(String key) {
        return args.get(key);
    }

    /*
     * return "okay" or "FAILED" based on equality of
     * the argument strings
     */
    private String cmp(String s1, String s2) {
        return s1.equals(s2) ? okay : failed;
    }

    private String cmp(String s1, int s2) {
        return cmp(s1, Integer.toString(s2));
    }

    public void run(boolean printHeader, boolean testRegression) {
        try {
            SSLHandshakeCompletedListener listener = null;

            if (printHeader)
                results.println(htmlHeader);
            results.println("SSL Client Tester");
            results.println(
                    "$Id$ " +
                            versionStr);

            String hostname;
            int port;

            String filename = getArgument("filename");

            if (isInvalid(filename)) {
                filename = "/index.html";
            }

            String msg = "GET " + filename;

            String portstr = getArgument("port");

            if (isInvalid(portstr)) {
                port = 443;
            } else {
                port = Integer.parseInt(portstr);
            }

            String addrstr = getArgument("ipaddr");
            hostname = addrstr; // unless it gets changed

            String tmpStr = getArgument("clientauth");
            if (isInvalid(tmpStr))
                doClientAuth = false;
            else {
                tmpStr = tmpStr.toLowerCase();
                doClientAuth = !(tmpStr.equals("off") ||
                        tmpStr.equals("false") ||
                        tmpStr.equals("0"));
            }

            if (isInvalid(addrstr)) {
                // check for a host name
                hostname = getArgument("hostname");
                if (isInvalid(hostname)) {
                    throw new Exception("hostname not specified");
                }
            }
            results.println("Connecting to " + hostname +
                    " on port " + port);

            SSLCertificateApprovalCallback approvalCallback = new TestCertApprovalCallback();
            SSLClientCertificateSelectionCallback certSelectionCallback = new TestClientCertificateSelectionCallback();
            Socket js = new Socket(InetAddress.getByName(hostname), port);
            //s = new SSLSocket(hostname, port, null, 0,
            try(SSLSocket s = new SSLSocket(js, hostname,
                    approvalCallback,
                    certSelectionCallback)) {

                s.forceHandshake();
                results.println("Connected.");

                // select the cert for client auth
                // You will have to provide a certificate with this
                // name if you expect client auth to work.

                //s.setClientCertNickname("JavaSSLTestClientCert");

                // Setup a handshake callback. This listener will get invoked
                // When the SSL handshake is completed on this socket.

                listener = new ClientHandshakeCB(this);
                s.addHandshakeCompletedListener(listener);

                //s.forceHandshake();

                OutputStream o = s.getOutputStream();

                PrintOutputStreamWriter out = new PrintOutputStreamWriter(o);

                results.println("Sending: " + msg + " to " + hostname +
                        ", " + port);

                // send HTTP GET message across SSL connection
                out.println(msg + "\r");

                InputStream in = s.getInputStream();
                byte[] bytes = new byte[4096];
                int totalBytes = 0;
                int numReads = 0;
                String lastBytes = null;

                // now try to read data back from the SSL connection
                try {
                    for (;;) {
                        results.println("Calling Read.");
                        int n = in.read(bytes, 0, bytes.length);
                        if (n == -1) {
                            results.println("EOF found.");
                            break;
                        }

                        if (n == 0) {
                            results.println("Zero bytes read?");
                            break;
                        }
                        numReads++;

                        if (totalBytes == 0) {
                            // don't print forever...
                            String data = new String(bytes, 0, 30, StandardCharsets.ISO_8859_1);
                            results.println("Read " + n + " bytes of data");
                            results.println("First 30 bytes: " + escapeHTML(data));
                        }

                        totalBytes += n;
                        lastBytes = new String(bytes, n - 31, 30, StandardCharsets.ISO_8859_1);
                    }
                } catch (IOException e) {
                    results.println(
                            "IOException while reading from pipe?  Actually got " +
                                    totalBytes + " bytes total");
                    e.printStackTrace(results);
                    results.println("");
                    throw e;
                } finally {
                    results.println("Last 30 bytes: " + lastBytes);
                    results.println("Number of read() calls: " + numReads);

                    /*
                     * if you want to test sslimpl.c's nsn_ThrowError(), try
                     * uncommenting the following line.  This will cause the
                     * getStatus() call to fail.
                     */

                    // in.close();

                    results.println("Diagnostics");
                    String tmp;
                    SSLSecurityStatus status = s.getStatus();

                    results.println("Total bytes read: " + totalBytes);
                    results.println("Security status of session:");
                    results.println(status.toString());

                    // now, for the regression testing stuff

                    if (testRegression) {
                        results.println("Regression Tests");

                        results.println("Handshake callback event happened: " +
                                ((handshakeEventHappened) ? okay : failed));

                        tmp = getArgument("filesize");
                        if (!isInvalid(tmp)) {
                            results.println("filesize: " + cmp(tmp, totalBytes));
                        }
                        tmp = getArgument("status");
                        if (!isInvalid(tmp)) {
                            results.println("status: " +
                                    cmp(tmp, status.getSecurityStatus()));
                        }
                        tmp = getArgument("sessionKeySize");
                        if (!isInvalid(tmp)) {
                            results.println("sessionKeySize: " +
                                    cmp(tmp, status.getSessionKeySize()));
                        }
                        tmp = getArgument("sessionSecretSize");
                        if (!isInvalid(tmp)) {
                            results.println("sessionSecretSize: " +
                                    cmp(tmp, status.getSessionSecretSize()));
                        }
                        tmp = getArgument("cipher");
                        if (!isInvalid(tmp)) {
                            results.println("cipher: " +
                                    cmp(tmp, status.getCipher()));
                        }
                        tmp = getArgument("issuer");
                        if (!isInvalid(tmp)) {
                            results.println("issuer: " +
                                    cmp(tmp, status.getRemoteIssuer()));
                        }
                        tmp = getArgument("subject");
                        if (!isInvalid(tmp)) {
                            results.println("subject: " +
                                    cmp(tmp, status.getRemoteSubject()));
                        }
                        tmp = getArgument("certSerialNum");
                        if (!isInvalid(tmp)) {
                            String serialNum = status.getSerialNumber();
                            results.println("certSerialNum: " +
                                    cmp(tmp, serialNum));
                        }
                    } // if false
                }
                // Got here, so no exception thrown above.
                // Try to clean up.
                o.close();
                o = null;
                in.close();
                in = null;
                if (listener != null) {
                    s.removeHandshakeCompletedListener(listener);
                    listener = null;
                }
            }

        } catch (Exception e) {
            results.println("***** TEST FAILED *****");
            e.printStackTrace(results);
            results.println("If there is no stack trace, try disabling the JIT and trying again.");
        }
        results.println("END OF TEST");
    }

    /**
     * given an input string, convert less-than, greater-than, and ampersand
     * from raw characters to escaped characters
     * (&lt; becomes `&amp;lt;', etc.)
     */
    private String escapeHTML(String s) {
        StringBuilder result = new StringBuilder();

        // this is inefficient, but I don't care
        for (int i = 0; i < s.length(); i++) {
            char c = s.charAt(i);
            switch (c) {
            case '<':
                result.append("&lt;");
                break;
            case '>':
                result.append("&gt;");
                break;
            case '&':
                result.append("&amp;");
                break;
            default:
                result.append(c);
                break;
            }
        }

        return result.toString();
    }

    public SSLClient(PrintStream ps, String verStr, String[] argv) {
        this.args = new HashMap<>();
        this.results = ps;
        this.versionStr = verStr;

        for (int i = 0; i < argNames.length; i++) {
            String value = values[i];
            if (value != null)
                this.args.put(argNames[i], value);
        }
        for (int i = 0; i < argv.length; i += 2) {
            this.args.put(argv[i], argv[i + 1]);
        }
    }

    static final int cipherSuites[] = {
            SSLSocket.SSL3_RSA_WITH_RC4_128_MD5,
            SSLSocket.SSL3_RSA_WITH_3DES_EDE_CBC_SHA,
            SSLSocket.SSL3_RSA_WITH_DES_CBC_SHA,
            SSLSocket.SSL3_RSA_EXPORT_WITH_RC4_40_MD5,
            SSLSocket.SSL3_RSA_EXPORT_WITH_RC2_CBC_40_MD5,
            SSLSocket.SSL3_RSA_WITH_NULL_MD5,
            0
    };

    public static void main(String[] argv) throws Exception {
        int i;

        try {
            InitializationValues vals = new InitializationValues(".");
            CryptoManager.initialize(vals);

            //	NSSInit.initialize("secmod.db", "key3.db", "cert7.db");
        } catch (KeyDatabaseException kdbe) {
            System.out.println("Couldn't open the key database");
            return;
        } catch (CertDatabaseException cdbe) {
            System.out.println("Couldn't open the certificate database");
            return;
        } catch (AlreadyInitializedException aie) {
            System.out.println("CryptoManager already initialized???");
            return;
        } catch (Exception e) {
            System.out.println("Exception occurred: " + e.getMessage());
            return;
        }

        /* enable all the SSL2 cipher suites */
        for (i = SSLSocket.SSL2_RC4_128_WITH_MD5; i <= SSLSocket.SSL2_DES_192_EDE3_CBC_WITH_MD5; ++i) {
            //	SSLSocket.setPermittedByPolicy(i, SSLSocket.SSL_ALLOWED);
            if (i != 0xFF05) {
                SSLSocket.setCipherPreferenceDefault(i, true);
            }
        }

        /* enable all the SSL3 cipher suites */
        for (i = 0; cipherSuites[i] != 0; ++i) {
            //	SSLSocket.setPermittedByPolicy(cipherSuites[i], SSLSocket.SSL_ALLOWED);
            SSLSocket.setCipherPreferenceDefault(cipherSuites[i], true);
        }

        SSLClient x = new SSLClient(System.out, "Stand alone Ver 0.01", argv);
        x.run(true, false);
    }

}

class ClientHandshakeCB implements SSLHandshakeCompletedListener {
    SSLClient sc;

    ClientHandshakeCB(SSLClient sc) {
        this.sc = sc;
    }

    @Override
    public void handshakeCompleted(SSLHandshakeCompletedEvent event) {
        try {
            sc.handshakeEventHappened = true;
            SSLSecurityStatus status = event.getStatus();
            sc.results.println("handshake happened\n" + status);
            System.out.println("Cert is " + status.getPeerCertificate());
        } catch (Exception e) {
            e.printStackTrace();
        }
    }
}
